import torch
import torchvision
from PIL import Image
import torch.nn.functional as F
from torch import nn

image_path = "../MNIST_dataset/MNIST/raw/test/11.jpg"  # 图片路径
image = Image.open(image_path)  # 将图片打开，此刻图片格式为PIL Image
# print(image)
transform = torchvision.transforms.Compose([torchvision.transforms.Resize((28, 28)), torchvision.transforms.ToTensor()])
image = transform(image)  # 将图片裁剪成28*28大小，然后将PIL Image格式转成Tensor类型
# print(image.shape)
image = image.cuda()  # 训练时用的是GPU，预测时也要用GPU 不然程序会报错


class MLP(nn.Module):  # 因为训练文件中保存模型用的是第一种保存方式，即：保存了模型和参数，所以在预测时，需要将构建的模型再写一遍。如果用的第二种保存方式，即：只保存参数文件，则预测时，不需要再次写模型代码
    def __init__(self):
        super(MLP, self).__init__()
        self.linear1 = nn.Linear(784, 512)
        self.linear2 = nn.Linear(512, 128)
        self.linear3 = nn.Linear(128, 10)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = F.relu(self.linear1(x))
        x = F.relu(self.linear2(x))
        x = F.softmax((self.linear3(x)), dim=1)
        return x


test_model = torch.load("mlp_49.pth")  # 加载选用的模型文件
test_model = test_model.cuda()  # 模型使用GPU
test_model.eval()  # 开始验证
with torch.no_grad():  # 验证时不需要求解梯度
    image = torch.reshape(image, (1, 1, 28, 28))  # 将上面转化的Tensor类型的图片添加一个batch_size维度，这样才符合预测图片要求。有个问题，这一步为什么不写也能正常运行？？？
    output = test_model(image)  # 得到预测的结果，是一个1*10的矩阵
print("预测的结果为：{}".format(output.argmax(1).item()))  # 找出1*10矩阵中的最大值，即为预测的数字
